import asyncio
import json
import logging
from datetime import datetime
from typing import Dict, List, Optional, Set

import pymysql
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, HTTPException
from pydantic import BaseModel

# 从主应用程序导入数据库配置
from mysql import DB_CONFIG

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("websocket")

app = FastAPI()

# 消息模型
class Message(BaseModel):
    sender_id: int
    receiver_id: int
    content: str
    msg_type: str = "text"  # 默认为文本消息，可以扩展为图片、视频等
    created_at: Optional[datetime] = None
    
# 连接管理器
class ConnectionManager:
    def __init__(self):
        # 用户连接映射表 {user_id: WebSocket}
        self.active_connections: Dict[int, WebSocket] = {}
        # 用户在线状态 {user_id: bool}
        self.online_status: Dict[int, bool] = {}
        
    async def connect(self, user_id: int, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[user_id] = websocket
        self.online_status[user_id] = True
        logger.info(f"用户 {user_id} 已连接")
        
    def disconnect(self, user_id: int):
        if user_id in self.active_connections:
            del self.active_connections[user_id]
            self.online_status[user_id] = False
            logger.info(f"用户 {user_id} 已断开连接")
    
    async def send_personal_message(self, message: str, user_id: int):
        if user_id in self.active_connections:
            await self.active_connections[user_id].send_text(message)
            return True
        return False
    
    def is_online(self, user_id: int) -> bool:
        return self.online_status.get(user_id, False)

# 创建连接管理器实例
manager = ConnectionManager()

# 获取数据库连接
def get_conn():
    return pymysql.connect(**DB_CONFIG)

# 保存消息到数据库
async def save_message_to_db(message: Message, is_read: bool = False):
    conn = get_conn()
    try:
        with conn.cursor() as cursor:
            sql = """
            INSERT INTO chat_messages 
            (sender_id, receiver_id, content, msg_type, is_read, created_at) 
            VALUES (%s, %s, %s, %s, %s, NOW())
            """
            cursor.execute(
                sql, 
                (
                    message.sender_id, 
                    message.receiver_id, 
                    message.content, 
                    message.msg_type, 
                    is_read
                )
            )
            message_id = cursor.lastrowid
            conn.commit()
            
            # 更新聊天会话
            await update_chat_session(message.sender_id, message.receiver_id, message_id)
            
            return message_id
    except Exception as e:
        logger.error(f"保存消息到数据库时出错: {str(e)}")
        return None
    finally:
        conn.close()

# 获取用户未读消息
async def get_unread_messages(user_id: int) -> List[Dict]:
    conn = get_conn()
    try:
        with conn.cursor(pymysql.cursors.DictCursor) as cursor:
            sql = """
            SELECT * FROM chat_messages 
            WHERE receiver_id = %s AND is_read = 0
            ORDER BY created_at ASC
            """
            cursor.execute(sql, (user_id,))
            messages = cursor.fetchall()
            # 将未读消息标记为已读
            if messages:
                update_sql = """
                UPDATE chat_messages 
                SET is_read = 1 
                WHERE receiver_id = %s AND is_read = 0
                """
                cursor.execute(update_sql, (user_id,))
                conn.commit()
            return messages
    except Exception as e:
        logger.error(f"获取未读消息时出错: {str(e)}")
        return []
    finally:
        conn.close()

# 获取用户聊天历史
async def get_chat_history(user_id: int, other_user_id: int, limit: int = 50) -> List[Dict]:
    conn = get_conn()
    try:
        with conn.cursor(pymysql.cursors.DictCursor) as cursor:
            sql = """
            SELECT * FROM chat_messages 
            WHERE (sender_id = %s AND receiver_id = %s) 
               OR (sender_id = %s AND receiver_id = %s)
            ORDER BY created_at DESC
            LIMIT %s
            """
            cursor.execute(sql, (user_id, other_user_id, other_user_id, user_id, limit))
            messages = cursor.fetchall()
            return messages
    except Exception as e:
        logger.error(f"获取聊天历史时出错: {str(e)}")
        return []
    finally:
        conn.close()

# 添加：更新或创建聊天会话
async def update_chat_session(sender_id: int, receiver_id: int, message_id: int):
    conn = get_conn()
    try:
        with conn.cursor() as cursor:
            # 确保user1_id小于user2_id，保持一致性
            user1_id, user2_id = (sender_id, receiver_id) if sender_id < receiver_id else (receiver_id, sender_id)
            
            # 更新或创建会话
            sql = """
            INSERT INTO chat_sessions (user1_id, user2_id, last_message_id, updated_at, created_at)
            VALUES (%s, %s, %s, NOW(), NOW())
            ON DUPLICATE KEY UPDATE last_message_id = %s, updated_at = NOW()
            """
            cursor.execute(sql, (user1_id, user2_id, message_id, message_id))
            conn.commit()
            return cursor.lastrowid
    except Exception as e:
        logger.error(f"更新聊天会话时出错: {str(e)}")
        return None
    finally:
        conn.close()

# WebSocket端点
@app.websocket("/ws/{user_id}")
async def websocket_endpoint(websocket: WebSocket, user_id: int):
    # 这里只需要发送者ID，因为一个用户可能会与多个人聊天
    # 接收者ID将在发送消息时指定
    await manager.connect(user_id, websocket)
    
    try:
        # 发送用户上线通知
        online_status = {
            "type": "status_update",
            "user_id": user_id,
            "status": "online",
            "timestamp": datetime.now().isoformat()
        }
        
        # 添加：更新数据库中的用户在线状态
        conn = get_conn()
        try:
            with conn.cursor() as cursor:
                sql = """
                INSERT INTO user_online_status (user_id, is_online, last_active_at)
                VALUES (%s, %s, NOW())
                ON DUPLICATE KEY UPDATE is_online = %s, last_active_at = NOW()
                """
                cursor.execute(sql, (user_id, True, True))
                conn.commit()
        except Exception as e:
            logger.error(f"更新用户在线状态时出错: {str(e)}")
        finally:
            conn.close()
        
        # 检查并发送未读消息
        unread_messages = await get_unread_messages(user_id)
        if unread_messages:
            unread_response = {
                "type": "unread_messages",
                "messages": unread_messages
            }
            await websocket.send_text(json.dumps(unread_response, default=str))
        
        # 接收和处理消息
        while True:
            data = await websocket.receive_text()
            try:
                message_data = json.loads(data)
                
                if message_data.get("type") == "chat_message":
                    # 创建消息对象
                    message = Message(
                        sender_id=user_id,
                        receiver_id=message_data.get("receiver_id"),
                        content=message_data.get("content"),
                        msg_type=message_data.get("msg_type", "text")
                    )
                    
                    # 准备发送的消息
                    message_json = {
                        "type": "chat_message",
                        "sender_id": message.sender_id,
                        "content": message.content,
                        "msg_type": message.msg_type,
                        "timestamp": datetime.now().isoformat()
                    }
                    
                    # 检查接收者是否在线
                    receiver_id = message.receiver_id
                    is_delivered = False
                    
                    if manager.is_online(receiver_id):
                        # 发送消息给接收者
                        is_delivered = await manager.send_personal_message(
                            json.dumps(message_json, default=str), 
                            receiver_id
                        )
                    
                    # 保存消息到数据库，如果已发送则标记为已读
                    await save_message_to_db(message, is_read=is_delivered)
                    
                    # 发送确认消息给发送者
                    confirmation = {
                        "type": "message_confirmation",
                        "message_id": message_data.get("message_id", "temp-" + str(id(message))),
                        "status": "delivered" if is_delivered else "sent",
                        "timestamp": datetime.now().isoformat()
                    }
                    await websocket.send_text(json.dumps(confirmation, default=str))
                
                elif message_data.get("type") == "fetch_history":
                    # 获取聊天历史
                    other_user_id = message_data.get("user_id")
                    limit = message_data.get("limit", 50)
                    
                    history = await get_chat_history(user_id, other_user_id, limit)
                    history_response = {
                        "type": "chat_history",
                        "user_id": other_user_id,
                        "messages": history
                    }
                    await websocket.send_text(json.dumps(history_response, default=str))
                
                elif message_data.get("type") == "typing_indicator":
                    # 发送正在输入指示器
                    receiver_id = message_data.get("receiver_id")
                    if manager.is_online(receiver_id):
                        typing_indicator = {
                            "type": "typing_indicator",
                            "user_id": user_id,
                            "status": message_data.get("status", "typing"),
                            "timestamp": datetime.now().isoformat()
                        }
                        await manager.send_personal_message(
                            json.dumps(typing_indicator, default=str),
                            receiver_id
                        )
            
            except json.JSONDecodeError:
                logger.error(f"无效的JSON数据: {data}")
            except Exception as e:
                logger.error(f"处理消息时出错: {str(e)}")
                
    except WebSocketDisconnect:
        manager.disconnect(user_id)
        # 发送用户离线通知
        offline_status = {
            "type": "status_update",
            "user_id": user_id,
            "status": "offline",
            "timestamp": datetime.now().isoformat()
        }
        # 广播用户离线状态（同样，可以根据需要修改为仅向好友发送）

# HTTP端点，用于检查用户是否在线
@app.get("/api/chat/user/{user_id}/status")
async def get_user_status(user_id: int):
    is_online = manager.is_online(user_id)
    return {
        "user_id": user_id,
        "status": "online" if is_online else "offline"
    }

# HTTP端点，获取用户未读消息数量
@app.get("/api/chat/user/{user_id}/unread_count")
async def get_unread_count(user_id: int):
    conn = get_conn()
    try:
        with conn.cursor() as cursor:
            sql = """
            SELECT COUNT(*) as count FROM chat_messages 
            WHERE receiver_id = %s AND is_read = 0
            """
            cursor.execute(sql, (user_id,))
            result = cursor.fetchone()
            return {
                "user_id": user_id,
                "unread_count": result[0] if result else 0
            }
    except Exception as e:
        logger.error(f"获取未读消息数量时出错: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        conn.close()

# 如果需要单独运行此WebSocket服务器
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("Websocket:app", host="0.0.0.0", port=8000, reload=True)